import socket
import json
from ollama import chat


def generate_response(messages):
    """调用Ollama生成回复"""
    try:
        response = chat('llama3.2', messages=messages)
        return str('\n'+response['message']['content']+'\n'+'——from llama'+'\n')
    except Exception as e:
        print(f"生成回复时出错: {e}")
        return "抱歉，生成回复时出错。"

def start_client():
    """启动TCP客户端"""
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client.connect(('localhost', 11435))
    print("客户端已连接到服务端")
    try:
        # 接收服务端初始消息
        server_message = client.recv(4096).decode()
        print(f"服务端: {server_message}")
        messages_S2C = [
            {
            'role': 'user',
            'content': server_message,
            },
        ]
        
        while True:
            client_response = generate_response(messages_S2C) 
            # client_response = generate_response(server_message)
            print(f"客户端回复: {client_response}")
            
            # 发送回复给服务端
            client.send(client_response.encode())
            
            # 接收服务端新消息
            server_message = client.recv(4096).decode()
            if not server_message:
                break
            print(f"服务端: {server_message}")
            
    except Exception as e:
        print(f"客户端通信时出错: {e}")
    finally:
        client.close()
        print("客户端已关闭")

if __name__ == "__main__":
    start_client()
